from utils import data

import numpy as np
import torch

"""
node_num, feat_dim, stat_dim, num_class, T
feat_Matrix, X_Node, X_Neis, dg_list
"""

content_path = "./data/cora/cora.content"
cite_path = "./data/cora/cora.cites"

# 读取文本内容
with open(content_path, "r") as fp:
    contents = fp.readlines()
with open(cite_path, "r") as fp:
    cites = fp.readlines()

contents = np.array([np.array(l.strip().split("\t")) for l in contents])
paper_list, feat_list, label_list = np.split(contents, [1,-1], axis= 1)
paper_list, label_list = np.squeeze(paper_list), np.squeeze(label_list)
# Paper -> Paper Index Dict
paper_dict = data.get_index_dict(paper_list)
# Label -> Label Index Dict
label_dict = data.get_index_dict(label_list)
labels = list(set(label_list))
# edge index
cites = [i.strip().split("\t") for i in cites]
cites = np.array([[paper_dict[i[0]], paper_dict[i[1]]] for i in cites], np.int64).T
cites = np.concatenate((cites, cites[::-1, :]), axis= 1)
# degree
_, degree_list = np.unique(cites[0,:], return_counts= True)

# input
node_num = len(paper_list)
feat_dim = feat_list.shape[1]
stat_fim = 32
num_class = len(labels)
T = 2
feat_Matrix = torch.Tensor(feat_list.astype(np.float32))
X_Node, X_Neis = np.split(cites, 2, axis= 0)
X_Node, X_Neis = torch.from_numpy(np.squeeze(X_Node)), torch.from_numpy(np.squeeze(X_Neis))
dg_list = degree_list[X_Node]
label_list = np.array([label_dict[i] for i in label_list])
label_list = torch.from_numpy(label_list)

if __name__ == "__main__":
    # print(paper_dict)
    # print(label_dict)
    print(cites.shape)
    print(torch.cuda.is_available())
    print(_)
    print(degree_list)